#include <stdio.h>
#include <string.h>
#include <time.h>
#include "crypto/sm2.h"

#define IO_BUF_SIZE 256
#define SM2_KEY_HEX_LEN 64

/**
 * Trim the last return char of a str.
 * @param str get rid of the return chars read by fgets,
 */
void line_trim(char *str) {
    for (int i = (int) strlen(str) - 1; i >= 0; i--) {
        if (str[i] == '\n' || str[i] == '\r') {
            str[i] = '\0';
            break;
        }
    }
}

/**
 * Sign the data provided by the context.
 *
 * @param input_file the input file contains the data provided on the provider
 * @param output_file the columns needed to fill the Excel
 */
void run_contest_sign(const char* input_file, const char* output_file) {
    FILE* input = fopen(input_file, "r");
    if (input == NULL) {
        printf("Failed to open file to read %s\n", input_file);
        return;
    }
    const int buf_size = IO_BUF_SIZE;
    char line_buf[IO_BUF_SIZE];
    bitstring* user_id = NULL;
    bigint* public_key_x = NULL;
    bigint* public_key_y = NULL;
    bigint* private_key = NULL;
    bitstring* msg = NULL;

    char* op = malloc(sizeof(char) *IO_BUF_SIZE);
    char* info = malloc(sizeof(char) *IO_BUF_SIZE);
    char* data_buf = malloc(sizeof(char) *IO_BUF_SIZE);

    while (fgets(line_buf, buf_size, input) != NULL) {
        sscanf(line_buf, "%s %s", op, info);
        if (strcmp(op, "SIGN") == 0) {
            printf("%s %s\n", op, info);
        } else if (strcmp(op, "USERID") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            user_id = bitstring_new_from_cstr(data_buf);
        } else if (strcmp(op, "PRIVATE_KEY") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            private_key = bigint_new_from_hex_str(data_buf);
        } else if (strcmp(op, "PUBLIC_KEY") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            // Parse public key
            const int len = (int) strlen(data_buf);
            const int offset = data_buf[0] == 0x04; // 0x04|X|Y| format
            if (len - offset != SM2_KEY_HEX_LEN * 2) {
                printf("ERROR: Public key length not correct\n");
                exit(1);
            }
            char key_buf[SM2_KEY_HEX_LEN + 1];
            key_buf[SM2_KEY_HEX_LEN] = '\0';
            strncpy(key_buf, data_buf + offset, SM2_KEY_HEX_LEN);
            public_key_x = bigint_new_from_hex_str(key_buf);
            strncpy(key_buf, data_buf + offset + SM2_KEY_HEX_LEN, SM2_KEY_HEX_LEN);
            public_key_y = bigint_new_from_hex_str(key_buf);
        } else if (strcmp(op, "MESSAGE") == 0) {
            int msg_max_len = buf_size;
            int msg_len = 0;
            char *msg_buf = malloc(sizeof(char) * buf_size);
            for (;;) {
                char ch = fgetc(input);
                if (msg_len == msg_max_len) { // realloc
                    msg_max_len += buf_size;
                    msg_buf = realloc(msg_buf, sizeof(char) * msg_max_len);
                    if (msg_buf == NULL) {
                        printf("Failed to allocate memory for message buffer\n");
                        exit(1);
                    }
                }
                if (ch == '\n' || ch == '\r' || ch == EOF) {
                    msg_buf[msg_len] = '\0';
                    // printf("Message size = %d\n%s\n", strlen(msg_buf), msg_buf);
                    msg = bitstring_new_from_hex_str(msg_buf);
                    break;
                }
                msg_buf[msg_len++] = ch;
            }
            free(msg_buf);
            break;
        }
    }
    free(op);
    free(info);
    free(data_buf);
    fclose(input);

    // Use default SM2 context
    sm2_ctx* ctx = sm2_ctx_new();
    sm2_ctx_init(ctx);
    sm2_ctx_set_private_key(ctx, private_key);
    sm2_ctx_set_public_key(ctx, public_key_x, public_key_y);

    // show_sm2_ctx(ctx);

    printf("* User Id hex\n");
    show_bitstring(user_id);
    bitstring* za = sm2_new_za(ctx, user_id);
    printf("* Za hex\n");
    show_bitstring(za);

    // Choose K
    bigint* k = bigint_new_from_hex_str(
        "6CB28D99385C175C94F94E934817663FC176D925DD72B727260DBAAE1FB2F988"
    );
    bigint_norm(k, ctx->n);
    printf("* k\n");
    show_bigint(k);


    FILE* output = fopen(output_file, "w");
    if (output == NULL) {
        printf("Failed to open file to write: %s\n", input_file);
        return;
    }
    fprintf(output, "RESULT contest\n");
    fprintf(output, "PUBLIC_KEY hex\n");
    write_bigint(output, ctx->Px);
    write_bigint(output, ctx->Py);
    clock_t sign_measure_start = clock();
    // Calculate digest
    bitstring* digest = sm2_digest(za, msg);
    printf("* Digest\n");
    show_bitstring(digest);
    fprintf(output, "DIGEST hex\n");
    write_bitstring(output, digest);

    // Sign digest with k
    sm2_sign_result* s_res = sm2_sign_with_k(ctx, digest, k);
    if (s_res->pass) {
        printf("* Sign Pass\n");
        printf("  r = ");
        show_bitstring(s_res->signature->r);
        printf("  s = ");
        show_bitstring(s_res->signature->s);
        printf("* ASN.1 in DER =\n");
        bitstring* der = sm2_signature_to_der(s_res->signature);
        show_bitstring(der);
        fprintf(output, "DER hex\n");
        write_bitstring(output, der);
        fprintf(output, "R hex\n");
        write_bitstring(output, s_res->signature->r);
        fprintf(output, "S hex\n");
        write_bitstring(output, s_res->signature->s);
        bitstring_delete(der);
    } else {
        printf("* Sign Fail: %s\n", s_res->error_message);
    }
    sm2_sign_result_delete(s_res);
    double benchmark = (double)(clock() - sign_measure_start) / CLOCKS_PER_SEC;
    printf("\n=================== Sign used %f sec ", benchmark);
    printf("=====================\n\n");
    bitstring_delete(digest);
    bigint_delete(k);
    sm2_ctx_delete(ctx);
    fclose(output);
}

/**
 * Verify the result generated.
 *
 * @param result_file the result file created by run_context_sign()
 */
void run_contest_verify(const char* result_file) {
    FILE* input = fopen(result_file, "r");
    if (input == NULL) {
        printf("Failed to open file to verify %s\n", result_file);
        return;
    }

    bigint* public_key_x = NULL;
    bigint* public_key_y = NULL;
    bitstring* digest = NULL;
    bitstring* r = NULL;
    bitstring* s = NULL;

    const int buf_size = IO_BUF_SIZE;
    char line_buf[IO_BUF_SIZE];
    char* op = malloc(sizeof(char) *IO_BUF_SIZE);
    char* info = malloc(sizeof(char) *IO_BUF_SIZE);
    char* data_buf = malloc(sizeof(char) *IO_BUF_SIZE);
    while (fgets(line_buf, buf_size, input) != NULL) {
        sscanf(line_buf, "%s %s", op, info);
        if (strcmp("RESULT", op) == 0) {
            continue;
        }
        if (strcmp(op, "PUBLIC_KEY") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            public_key_x = bigint_new_from_hex_str(data_buf);
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            public_key_y = bigint_new_from_hex_str(data_buf);
        } else if (strcmp(op, "DIGEST") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            digest = bitstring_new_from_hex_str(data_buf);
        } else if (strcmp(op, "R") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            r = bitstring_new_from_hex_str(data_buf);
        } else if (strcmp(op, "S") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            s = bitstring_new_from_hex_str(data_buf);
        } else if (strcmp(op, "DER") == 0) {
            fgets(data_buf, buf_size, input);
            // ignored.
        }
    }

    clock_t v_measure_start = clock();
    sm2_ctx* ctx = sm2_ctx_new();
    sm2_ctx_init(ctx);
    sm2_ctx_set_public_key(ctx, public_key_x, public_key_y);

    // show_sm2_ctx(ctx);
    printf("* Verification Starts \n");
    printf("  h = ");
    show_bitstring(digest);
    printf("  r = ");
    show_bitstring(r);
    printf("  s = ");
    show_bitstring(s);
    sm2_signature signature = { .r = r, .s = s };
    sm2_verify_result* v_res = sm2_verify(
        ctx,
        digest,
        &signature
    );
    if (v_res->pass) {
        printf("[√] Verification passed\n");
    } else {
        printf("[x] Verification failed: %s\n", v_res->error_message);
    }
    double benchmark = (double)(clock() - v_measure_start) / CLOCKS_PER_SEC;
    printf("\n=================== Verify used %f sec ", benchmark);
    printf("=====================\n\n");

    bigint_delete(public_key_x);
    bigint_delete(public_key_y);
    bitstring_delete(digest);
    bitstring_delete(r);
    bitstring_delete(s);
    sm2_ctx_delete(ctx);
    free(op);
    free(info);
    free(data_buf);
    fclose(input);
}

int main(void) {
    // initializes the random number generator with the current time,
    // srand((unsigned int)time(NULL));
    // test_sm3_hash_2();
    // test_sm3_hash_2();
    // test_bigint_1();
    // test_sm2_ecc();
    // test_sm2_sign();
    // test_sm2_verify();

    run_contest_sign(
        "../data/contest_case_1.txt",
        "../data/contest_case_1_out.txt"
    );
    run_contest_verify(
        "../data/contest_case_1_out.txt"
    );
    run_contest_sign(
        "../data/contest_case_2.txt",
        "../data/contest_case_2_out.txt"
    );
    run_contest_verify(
        "../data/contest_case_2_out.txt"
    );
    run_contest_sign(
        "../data/contest_case_3.txt",
        "../data/contest_case_3_out.txt"
    );
    run_contest_verify(
        "../data/contest_case_3_out.txt"
    );
    printf("# All task finished\n");
    return 0;
}